import os

import tensorflow as tf
from keras.datasets import cifar10
from tensorflow.python import keras
from tensorflow.python.keras import layers, optimizers

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'


def preprocess(x, y):
    # [0,256]--->[-1,1]
    x = tf.cast(x, dtype=tf.float32) / 255 - 1
    y = tf.cast(y, tf.int32)
    return x, y


batchsz = 128

(x, y), (x_val, y_val) = cifar10.load_data()
y = tf.squeeze(y)
y_val = tf.squeeze(y_val)
y = tf.one_hot(y, depth=10)  # [50k,10]
y_val = tf.one_hot(y_val, depth=10)
print("datasets", x.shape, y.shape, x_val.shape, y_val.shape, x.min(), x.max())

train_db = tf.data.Dataset.from_tensor_slices((x, y))
train_db = train_db.map(preprocess).shuffle(10000).batch(batchsz)

test_db = tf.data.Dataset.from_tensor_slices((x, y))
test_db = test_db.map(preprocess).batch(batchsz)

sample = next(iter(train_db))
print("batch shape", sample[0].shape, sample[1].shape)


# 自定义层
class MyDense(layers.Layer):
    def __init__(self, inp_dim, oup_dim):
        super(MyDense, self).__init__()
        self.kernel = self.add_weight('w', [inp_dim, oup_dim])
        # self.bias = self.add_variable('b', [oup_dim])

    def call(self, inputs, *args, **kwargs):
        x = inputs @ self.kernel
        # x = inputs @ self.kernel + self.bias
        return x

    # def acll(self, inputs, *args, **kwargs):
    #     out = inputs @ self.kernel
    #     return out


# 自定义网络
class MyNetwork(keras.Model):
    def __init__(self):
        super(MyNetwork, self).__init__()
        self.fc1 = MyDense(32 * 32 * 3, 256)
        self.fc2 = MyDense(256, 128)
        self.fc3 = MyDense(128, 64)
        self.fc4 = MyDense(64, 32)
        self.fc5 = MyDense(32, 10)

    def call(self, inputs, training=None, mask=None):
        x = tf.reshape(inputs, [-1, 32 * 32 * 3])
        x = self.fc1(x)  # [b,32*32*3]------>[b,256]
        x = tf.nn.relu(x)
        x = self.fc2(x)  # [b,256]------>[b,128]
        x = tf.nn.relu(x)
        x = self.fc3(x)  # [b,128]------>[b,64]
        x = tf.nn.relu(x)
        x = self.fc4(x)  # [b,64]------>[b,32]
        x = tf.nn.relu(x)
        x = self.fc5(x)  # [b,32]------>[b,10]
        return x


network = MyNetwork()
network.compile(optimizer=optimizers.adam_v2.Adam(learning_rate=1e-3),
                loss=tf.losses.CategoricalCrossentropy(from_logits=True),  # from_logits=True表示应用softmax
                metrics=['accuracy'])

network.fit(train_db, epochs=10, validation_data=test_db, validation_freq=2)
network.evaluate(test_db)

network.save_weights('ljh/Myl')
print("succ save weight")

del network
print("succ del net")
ljh = MyNetwork()

# 加载网络结构
ljh.compile(optimizer=optimizers.adam_v2.Adam(learning_rate=1e-3),
            loss=tf.losses.CategoricalCrossentropy(from_logits=True),
            metrics=['accuracy'])


ljh.load_weights('ljh/Myl')
print("succ load weight")
ljh.evaluate(test_db)
